from nnmnkwii.datasets import FileDataSource, FileSourceDataset
from nnmnkwii.datasets import MemoryCacheFramewiseDataset
from nnmnkwii.preprocessing import trim_zeros_frames, remove_zeros_frames
from nnmnkwii.preprocessing import minmax, meanvar, minmax_scale, scale
from nnmnkwii import paramgen
from nnmnkwii.io import hts
from nnmnkwii.frontend import merlin as fe
from nnmnkwii.postfilters import merlin_post_filter

from os.path import join, expanduser, basename, splitext, basename, exists
import os
from glob import glob
import numpy as np
from scipy.io import wavfile
from sklearn.model_selection import train_test_split
import pyworld
import pysptk
import librosa
import librosa.display
import IPython
from IPython.display import Audio
import config as cfg


class BinaryFileSource(FileDataSource):
    def __init__(self, data_root, dim, train):
        self.data_root = data_root
        self.dim = dim
        self.train = train

    def collect_files(self):
        files = sorted(glob(join(self.data_root, "*.bin")))
        files = files[:len(files) - 5]  # last 5 is real testset
        train_files, test_files = train_test_split(files,
                                                   test_size=cfg.test_size,
                                                   random_state=cfg.random_state)
        if self.train:
            return train_files
        else:
            return test_files

    def collect_features(self, path):
        return np.fromfile(path, dtype=np.float32).reshape(-1, self.dim)


X = {"duration": {}, "acoustic": {}}
Y = {"duration": {}, "acoustic": {}}
utt_lengths = {"duration": {}, "acoustic": {}}
for ty in ["duration", "acoustic"]:
    for phase in ["train", "test"]:
        train = phase == "train"
        x_dim = cfg.duration_linguistic_dim if ty == "duration" else cfg.acoustic_linguisic_dim
        y_dim = cfg.duration_dim if ty == "duration" else cfg.acoustic_dim
        X[ty][phase] = FileSourceDataset(
            BinaryFileSource(join(cfg.DATA_ROOT, "X_{}".format(ty)),
                             dim=x_dim,
                             train=train))
        Y[ty][phase] = FileSourceDataset(
            BinaryFileSource(join(cfg.DATA_ROOT, "Y_{}".format(ty)),
                             dim=y_dim,
                             train=train))
        # this triggers file loads, but can be neglectable in terms of performance.
        utt_lengths[ty][phase] = [len(x) for x in X[ty][phase]]


X_min = {}
X_max = {}
Y_mean = {}
Y_var = {}
Y_scale = {}

for typ in ["acoustic", "duration"]:
    X_min[typ], X_max[typ] = minmax(X[typ]["train"], utt_lengths[typ]["train"])
    Y_mean[typ], Y_var[typ] = meanvar(Y[typ]["train"], utt_lengths[typ]["train"])
    Y_scale[typ] = np.sqrt(Y_var[typ])


# print(X_min)

X_dur_train = X['duration']['train']
Y_dur_train = Y['duration']['train']
X_min_dur = X_min['duration']
X_max_dur = X_max['duration']
print(X_min_dur.shape)
print(X_max_dur.shape)
print()
for ele in Y_dur_train:
    print(ele[1])
    # minmax_scale(ele, X_min_dur, X_max_dur)
